import cv2
import pyzed.sl as sl
import math
import numpy as np
zed = sl.Camera()
# Create configuration parameters
init_params = sl.InitParameters()
init_params.sdk_verbose = True # Enable the verbose modeinit_params.depth_mode = sl.DEPTH_MODE.PERFORMANCE 
init_params.camera_resolution = sl.RESOLUTION.HD720 # Use HD1080 video mode
init_params.camera_fps = 30 # Set fps at 30
# Set the depth mode to performance (fastest)
# Open the camera
err = zed.open(init_params)
if (err!=sl.ERROR_CODE.SUCCESS):
    exit(-1)

class zedCameraFrame():
    def __init__(self):
        self.left = sl.Mat()
        self.right = sl.Mat()
        self.depth = sl.Mat()
        self.point_cloud = sl.Mat()
        if (zed.grab() == sl.ERROR_CODE.SUCCESS) :
            self.state = True
            zed.retrieve_image(self.left, sl.VIEW.LEFT)
            zed.retrieve_image(self.right, sl.VIEW.RIGHT)
            zed.retrieve_measure(self.depth, sl.MEASURE.DEPTH) 
            zed.retrieve_measure(self.point_cloud, sl.MEASURE.XYZRGBA)
        else:
            self.state = False
    
    def get_img(self, pic_left=True):
        '''
        pic_left == True: use left pic
        pic_left == False: use right pic
        '''
        if self.state:
            if pic_left:
                img = self.left.get_data()
            else:
                img = self.right.get_data()
            b = np.zeros((img.shape[0],img.shape[1]),dtype=img.dtype)
            g = np.zeros((img.shape[0],img.shape[1]),dtype=img.dtype)
            r = np.zeros((img.shape[0],img.shape[1]),dtype=img.dtype)

            b[:,:] = img[:,:,0]  # 复制 b 通道的数据
            g[:,:] = img[:,:,1]  # 复制 g 通道的数据
            r[:,:] = img[:,:,2]  # 复制 r 通道的数据
            return [True,cv2.merge([b,g,r])]
        return [False,None]

    def get_depth(self,points):
        distance = []
        for x,y in points:
            point_cloud_value = self.point_cloud.get_value(x, y)[1]
            distance.append(math.sqrt(point_cloud_value[0]*point_cloud_value[0] + point_cloud_value[1]*point_cloud_value[1] + point_cloud_value[2]*point_cloud_value[2]))
        return distance

def get_depth(points):
    # Capture 50 images and depth, then stop
    image = sl.Mat()
    depth = sl.Mat()
    point_cloud = sl.Mat()
    # Grab an image    
    if (zed.grab() == sl.ERROR_CODE.SUCCESS) :
        # A new image is available if grab() returns SUCCESS        
        zed.retrieve_image(image, sl.VIEW.LEFT) 
        # Get the left image        
        zed.retrieve_measure(depth, sl.MEASURE.DEPTH) 
        # Retrieve depth Mat. Depth is aligned on the left image
        zed.retrieve_measure(point_cloud, sl.MEASURE.XYZRGBA)          
        # Retrieve colored point cloud. Point cloud is aligned on the left image 
        distance = []
        for x,y in points:
            point_cloud_value = point_cloud.get_value(x, y)[1]
            distance.append(math.sqrt(point_cloud_value[0]*point_cloud_value[0] + point_cloud_value[1]*point_cloud_value[1] + point_cloud_value[2]*point_cloud_value[2]))
        return distance

def get_img():
    left_img = sl.Mat()
    runtimeParas = sl.RuntimeParameters() 
    if zed.grab(runtimeParas) == sl.ERROR_CODE.SUCCESS:
        zed.retrieve_image(left_img, view=sl.VIEW.LEFT)
        #r,g,b = cv2.split(left_img.get_data())
        img = left_img.get_data()
        b = np.zeros((img.shape[0],img.shape[1]),dtype=img.dtype)
        g = np.zeros((img.shape[0],img.shape[1]),dtype=img.dtype)
        r = np.zeros((img.shape[0],img.shape[1]),dtype=img.dtype)

        b[:,:] = img[:,:,0]  # 复制 b 通道的数据
        g[:,:] = img[:,:,1]  # 复制 g 通道的数据
        r[:,:] = img[:,:,2]  # 复制 r 通道的数据
        return [True,cv2.merge([b,g,r])]
    return [False,None]

def __main():
    p = True
    while True:
        '''ret, img = get_img()
        cv2.imshow('img',img)
        cv2.waitKey(2)'''
        #print(get_depth([[400,400],[800,800]]))
        p = not p
        frame = zedCameraFrame()
        _,a = frame.get_img()
        cv2.imshow('img',a)
        cv2.waitKey(2)
        print(frame.get_depth([[400,400],[800,800]]))
    
if __name__ == '__main__':
    __main()